// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <melon/random.h>
#include <melon/init/init.h>

#include <pollux/testing/gtest_utils.h>
#include <pollux/testing/exec/util/operator_test_base.h>
#include <pollux/plan/plan_builder.h>
#include <pollux/vector/vector_maker.h>

#include <pollux/substrait/substrait_to_pollux_plan.h>
#include <pollux/substrait/pollux_to_substrait_plan.h>

#include <pollux/substrait/variant_to_vector_converter.h>

using namespace kumo::pollux;
using namespace kumo::pollux::test;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::exec::test;
using namespace kumo::pollux::substrait;
using namespace kumo::pollux::plan;

class PolluxSubstraitRoundTripTest : public OperatorTestBase {
protected:
    /// Makes a vector of INTEGER type with 'size' RowVectorPtr.
    /// @param size The number of RowVectorPtr.
    /// @param childSize The number of columns for each row.
    /// @param batchSize The batch Size of the data.
    std::vector<RowVectorPtr>
    makeVectors(int64_t size, int64_t childSize, int64_t batchSize) {
        std::vector<RowVectorPtr> vectors;
        std::mt19937 gen(std::mt19937::default_seed);
        for (int i = 0; i < size; i++) {
            std::vector<VectorPtr> children;
            for (int j = 0; j < childSize; j++) {
                children.emplace_back(make_flat_vector<int32_t>(
                    batchSize,
                    [&](auto /*row*/) {
                        return melon::Random::rand32(INT32_MAX / 4, INT32_MAX / 2, gen);
                    },
                    null_every(2)));
            }

            vectors.push_back(make_row_vector({children}));
        }
        return vectors;
    }

    void assertPlanConversion(
        const std::shared_ptr<const core::PlanNode> &plan,
        const std::string &duckDbSql) {
        assertQuery(plan, duckDbSql);

        // Convert Pollux Plan to Substrait Plan.
        google::protobuf::Arena arena;
        auto substraitPlan = polluxConvertor_->toSubstrait(arena, plan);

        // Convert Substrait Plan to the same Pollux Plan.
        auto samePlan = substraitConverter_->toPolluxPlan(substraitPlan);

        // Assert pollux again.
        assertQuery(samePlan, duckDbSql);
    }

    void assertFailingPlanConversion(
        const std::shared_ptr<const core::PlanNode> &plan,
        const std::string &expectedErrorMessage) {
        CursorParameters params;
        params.planNode = plan;
        POLLUX_ASSERT_THROW(
            readCursor(params, [](auto /*task*/) {}), expectedErrorMessage);

        // Convert Pollux Plan to Substrait Plan.
        google::protobuf::Arena arena;
        auto substraitPlan = polluxConvertor_->toSubstrait(arena, plan);

        // Convert Substrait Plan to the same Pollux Plan.
        auto samePlan = substraitConverter_->toPolluxPlan(substraitPlan);

        // Assert pollux again.
        params.planNode = samePlan;
        POLLUX_ASSERT_THROW(
            readCursor(params, [](auto /*task*/) {}), expectedErrorMessage);
    }

    std::shared_ptr<PolluxToSubstraitPlanConvertor> polluxConvertor_ =
            std::make_shared<PolluxToSubstraitPlanConvertor>();
    std::shared_ptr<SubstraitPolluxPlanConverter> substraitConverter_ =
            std::make_shared<SubstraitPolluxPlanConverter>(pool_.get());
};

TEST_F(PolluxSubstraitRoundTripTest, project) {
    auto vectors = makeVectors(3, 4, 2);
    createDuckDbTable(vectors);
    auto plan =
            PlanBuilder().values(vectors).project({"c0 + c1", "c1 / c2"}).planNode();
    assertPlanConversion(plan, "SELECT c0 + c1, c1 // c2 FROM tmp");
}

TEST_F(PolluxSubstraitRoundTripTest, cast) {
    auto vectors = makeVectors(3, 4, 2);
    createDuckDbTable(vectors);
    // Cast int32 to int64.
    auto plan =
            PlanBuilder().values(vectors).project({"cast(c0 as bigint)"}).planNode();
    assertPlanConversion(plan, "SELECT cast(c0 as bigint) FROM tmp");

    // Cast literal "abc" to int64 and allow cast failure, expecting no exception.
    plan = PlanBuilder()
            .values(vectors)
            .project({"try_cast('abc' as bigint)"})
            .planNode();
    assertPlanConversion(plan, "SELECT try_cast('abc' as bigint) FROM tmp");

    // Cast literal "abc" to int64, expecting an exception to be thrown.
    plan = PlanBuilder()
            .values(vectors)
            .project({"cast('abc' as bigint)"})
            .planNode();
    assertFailingPlanConversion(plan, "Cannot cast VARCHAR 'abc' to BIGINT");
}

TEST_F(PolluxSubstraitRoundTripTest, filter) {
    auto vectors = makeVectors(3, 4, 2);
    createDuckDbTable(vectors);

    auto plan = PlanBuilder().values(vectors).filter("c2 < 1000").planNode();
    assertPlanConversion(plan, "SELECT * FROM tmp WHERE c2 < 1000");
}

TEST_F(PolluxSubstraitRoundTripTest, null) {
    auto vectors = make_row_vector(ROW({}, {}), 1);
    auto plan = PlanBuilder().values({vectors}).project({"NULL"}).planNode();
    assertPlanConversion(plan, "SELECT NULL ");
}

TEST_F(PolluxSubstraitRoundTripTest, values) {
    RowVectorPtr vectors = make_row_vector(
        {
            make_flat_vector<int64_t>(
                {2499109626526694126, 2342493223442167775, 4077358421272316858}),
            make_flat_vector<int32_t>({581869302, -708632711, -133711905}),
            make_flat_vector<double>(
                {0.90579193414549275, 0.96886777112423139, 0.63235925003444637}),
            make_flat_vector<bool>({true, false, false}),
            make_flat_vector<int32_t>(3, nullptr, null_every(1))

        });
    createDuckDbTable({vectors});

    auto plan = PlanBuilder().values({vectors}).planNode();

    assertPlanConversion(plan, "SELECT * FROM tmp");
}

TEST_F(PolluxSubstraitRoundTripTest, count) {
    auto vectors = makeVectors(2, 7, 3);
    createDuckDbTable(vectors);

    auto plan = PlanBuilder()
            .values(vectors)
            .filter("c6 < 24")
            .singleAggregation({"c0", "c1"}, {"count(c4) as num_price"})
            .project({"num_price"})
            .planNode();

    assertPlanConversion(
        plan,
        "SELECT count(c4) as num_price FROM tmp WHERE c6 < 24 GROUP BY c0, c1");
}

TEST_F(PolluxSubstraitRoundTripTest, countAll) {
    auto vectors = makeVectors(2, 7, 3);
    createDuckDbTable(vectors);

    auto plan = PlanBuilder()
            .values(vectors)
            .filter("c6 < 24")
            .singleAggregation({"c0", "c1"}, {"count(1) as num_price"})
            .project({"num_price"})
            .planNode();

    assertPlanConversion(
        plan,
        "SELECT count(*) as num_price FROM tmp WHERE c6 < 24 GROUP BY c0, c1");
}

TEST_F(PolluxSubstraitRoundTripTest, sum) {
    auto vectors = makeVectors(2, 7, 3);
    createDuckDbTable(vectors);

    auto plan = PlanBuilder()
            .values(vectors)
            .partialAggregation({}, {"sum(1)", "count(c4)"})
            .planNode();

    assertPlanConversion(plan, "SELECT sum(1), count(c4) FROM tmp");
}

TEST_F(PolluxSubstraitRoundTripTest, sumAndCount) {
    auto vectors = makeVectors(2, 7, 3);
    createDuckDbTable(vectors);

    auto plan = PlanBuilder()
            .values(vectors)
            .partialAggregation({}, {"sum(c1)", "count(c4)"})
            .finalAggregation()
            .planNode();

    assertPlanConversion(plan, "SELECT sum(c1), count(c4) FROM tmp");
}

TEST_F(PolluxSubstraitRoundTripTest, sumGlobal) {
    auto vectors = makeVectors(2, 7, 3);
    createDuckDbTable(vectors);

    // Global final aggregation.
    auto plan = PlanBuilder()
            .values(vectors)
            .partialAggregation({"c0"}, {"sum(c0)", "sum(c1)"})
            .intermediateAggregation()
            .finalAggregation()
            .planNode();
    assertPlanConversion(
        plan, "SELECT c0, sum(c0), sum(c1) FROM tmp GROUP BY c0");
}

TEST_F(PolluxSubstraitRoundTripTest, sumMask) {
    auto vectors = makeVectors(2, 7, 3);
    createDuckDbTable(vectors);

    auto plan =
            PlanBuilder()
            .values(vectors)
            .project({"c0", "c1", "c2 % 2 < 10 AS m0", "c3 % 3 = 0 AS m1"})
            .partialAggregation(
                {}, {"sum(c0)", "sum(c0)", "sum(c1)"}, {"m0", "m1", "m1"})
            .finalAggregation()
            .planNode();

    assertPlanConversion(
        plan,
        "SELECT sum(c0) FILTER (WHERE c2 % 2 < 10), "
        "sum(c0) FILTER (WHERE c3 % 3 = 0), sum(c1) FILTER (WHERE c3 % 3 = 0) "
        "FROM tmp");
}

TEST_F(PolluxSubstraitRoundTripTest, rowConstructor) {
    RowVectorPtr vectors = make_row_vector(
        {
            make_flat_vector<double_t>({0.905791934145, 0.968867771124}),
            make_flat_vector<int64_t>({2499109626526694126, 2342493223442167775}),
            make_flat_vector<int32_t>({581869302, -133711905})
        });
    createDuckDbTable({vectors});

    auto plan = PlanBuilder()
            .values({vectors})
            .project({"row_constructor(c1, c2)"})
            .planNode();
    assertPlanConversion(plan, "SELECT row(c1, c2) FROM tmp");
}

TEST_F(PolluxSubstraitRoundTripTest, projectAs) {
    RowVectorPtr vectors = make_row_vector(
        {
            make_flat_vector<double_t>({0.905791934145, 0.968867771124}),
            make_flat_vector<int64_t>({2499109626526694126, 2342493223442167775}),
            make_flat_vector<int32_t>({581869302, -133711905})
        });
    createDuckDbTable({vectors});

    auto plan = PlanBuilder()
            .values({vectors})
            .filter("c0 < 0.5")
            .project({"c1 * c2 as revenue"})
            .partialAggregation({}, {"sum(revenue)"})
            .planNode();
    assertPlanConversion(
        plan, "SELECT sum(c1 * c2) as revenue FROM tmp WHERE c0 < 0.5");
}

TEST_F(PolluxSubstraitRoundTripTest, avg) {
    auto vectors = makeVectors(2, 7, 3);
    createDuckDbTable(vectors);

    auto plan = PlanBuilder()
            .values(vectors)
            .partialAggregation({}, {"avg(c4)"})
            .finalAggregation()
            .planNode();

    assertPlanConversion(plan, "SELECT avg(c4) FROM tmp");
}

TEST_F(PolluxSubstraitRoundTripTest, caseWhen) {
    auto vectors = makeVectors(3, 4, 2);
    createDuckDbTable(vectors);
    auto plan =
            PlanBuilder()
            .values(vectors)
            .project(
                {"case when c0=1 then c1 when c0=2 then c2 else c3  end as x"})
            .planNode();

    assertPlanConversion(
        plan,
        "SELECT case when c0=1 then c1 when c0=2 then c2 else c3 end as x FROM tmp");

    // Switch expression without else.
    plan = PlanBuilder()
            .values(vectors)
            .project({"case when c0=1 then c1 when c0=2 then c2 end as x"})
            .planNode();
    assertPlanConversion(
        plan,
        "SELECT case when c0=1 then c1 when c0=2 then c2  end as x FROM tmp");
}

TEST_F(PolluxSubstraitRoundTripTest, ifThen) {
    auto vectors = makeVectors(3, 4, 2);
    createDuckDbTable(vectors);
    auto plan = PlanBuilder()
            .values(vectors)
            .project({"if (c0=1, c0 + 1, c1 + 2) as x"})
            .planNode();
    assertPlanConversion(plan, "SELECT if (c0=1, c0 + 1, c1 + 2) as x FROM tmp");
}

TEST_F(PolluxSubstraitRoundTripTest, orderBySingleKey) {
    auto vectors = makeVectors(10, 4, 2);
    createDuckDbTable(vectors);
    auto plan = PlanBuilder()
            .values(vectors)
            .orderBy({"c0 DESC NULLS LAST"}, false)
            .planNode();
    assertPlanConversion(plan, "SELECT * FROM tmp ORDER BY c0 DESC NULLS LAST");
}

TEST_F(PolluxSubstraitRoundTripTest, orderBy) {
    auto vectors = makeVectors(10, 4, 2);
    createDuckDbTable(vectors);
    auto plan = PlanBuilder()
            .values(vectors)
            .orderBy({"c0 ASC NULLS FIRST", "c1 ASC NULLS LAST"}, false)
            .planNode();
    assertPlanConversion(
        plan, "SELECT * FROM tmp ORDER BY c0 NULLS FIRST, c1 NULLS LAST");
}

TEST_F(PolluxSubstraitRoundTripTest, limit) {
    auto vectors = makeVectors(10, 4, 2);
    createDuckDbTable(vectors);
    auto plan = PlanBuilder().values(vectors).limit(0, 10, false).planNode();
    assertPlanConversion(plan, "SELECT * FROM tmp LIMIT 10");

    // With offset.
    plan = PlanBuilder().values(vectors).limit(5, 10, false).planNode();
    assertPlanConversion(plan, "SELECT * FROM tmp OFFSET 5 LIMIT 10");
}

TEST_F(PolluxSubstraitRoundTripTest, topN) {
    auto vectors = makeVectors(10, 4, 2);
    createDuckDbTable(vectors);
    auto plan = PlanBuilder()
            .values(vectors)
            .topN({"c0 NULLS FIRST"}, 10, false)
            .planNode();
    assertPlanConversion(
        plan, "SELECT * FROM tmp ORDER BY c0 NULLS FIRST LIMIT 10");
}

TEST_F(PolluxSubstraitRoundTripTest, topNFilter) {
    auto vectors = makeVectors(10, 4, 2);
    createDuckDbTable(vectors);
    auto plan = PlanBuilder()
            .values(vectors)
            .filter("c0 > 15")
            .topN({"c0 DESC NULLS FIRST"}, 10, false)
            .planNode();
    assertPlanConversion(
        plan,
        "SELECT * FROM tmp WHERE c0 > 15 ORDER BY c0 DESC NULLS FIRST LIMIT 10");
}

TEST_F(PolluxSubstraitRoundTripTest, topNTwoKeys) {
    auto vectors = makeVectors(10, 4, 2);
    createDuckDbTable(vectors);
    auto plan = PlanBuilder()
            .values(vectors)
            .filter("c0 > 15")
            .topN({"c0 NULLS FIRST", "c1 DESC NULLS LAST"}, 10, false)
            .planNode();
    assertPlanConversion(
        plan,
        "SELECT * FROM tmp WHERE c0 > 15 ORDER BY c0 NULLS FIRST, c1 DESC NULLS LAST LIMIT 10");
}

namespace {
    core::TypedExprPtr makeConstantExpr(const TypePtr &type, const variant &value) {
        return std::make_shared<const core::ConstantTypedExpr>(type, value);
    }

    core::TypedExprPtr makeConstantExpr(const VectorPtr &vector) {
        return std::make_shared<const core::ConstantTypedExpr>(
            BaseVector::wrap_in_constant(1, 0, vector));
    }
} // namespace

TEST_F(PolluxSubstraitRoundTripTest, notNullLiteral) {
    auto vectors = make_row_vector(ROW({}, {}), 1);
    auto plan = PlanBuilder(pool_.get())
            .values({vectors})
            .addNode([&](std::string id, core::PlanNodePtr input) {
                std::vector<std::string> projectNames = {
                    "a", "b", "c", "d", "e", "f", "g", "h"
                };
                std::vector<core::TypedExprPtr> projectExpressions = {
                    makeConstantExpr(BOOLEAN(), (bool) 1),
                    makeConstantExpr(TINYINT(), (int8_t) 23),
                    makeConstantExpr(SMALLINT(), (int16_t) 45),
                    makeConstantExpr(INTEGER(), (int32_t) 678),
                    makeConstantExpr(BIGINT(), (int64_t) 910),
                    makeConstantExpr(REAL(), (float) 1.23),
                    makeConstantExpr(DOUBLE(), (double) 4.56),
                    makeConstantExpr(VARCHAR(), "789")
                };
                return std::make_shared<core::ProjectNode>(
                    id,
                    std::move(projectNames),
                    std::move(projectExpressions),
                    input);
            })
            .planNode();
    assertPlanConversion(
        plan, "SELECT true, 23, 45, 678, 910, 1.23, 4.56, '789'");
}

TEST_F(PolluxSubstraitRoundTripTest, arrayLiteral) {
    auto vectors = make_row_vector(ROW({}), 1);
    auto plan =
            PlanBuilder(pool_.get())
            .values({vectors})
            .addNode([&](std::string id, core::PlanNodePtr input) {
                std::vector<core::TypedExprPtr> expressions = {
                    makeConstantExpr(
                        make_nullable_array_vector<bool>({{true, std::nullopt}})),
                    makeConstantExpr(
                        make_nullable_array_vector<int8_t>({{0, std::nullopt}})),
                    makeConstantExpr(
                        make_nullable_array_vector<int16_t>({{1, std::nullopt}})),
                    makeConstantExpr(
                        make_nullable_array_vector<int32_t>({{2, std::nullopt}})),
                    makeConstantExpr(
                        make_nullable_array_vector<int64_t>({{3, std::nullopt}})),
                    makeConstantExpr(
                        make_nullable_array_vector<float>({{4.4, std::nullopt}})),
                    makeConstantExpr(
                        make_nullable_array_vector<double>({{5.5, std::nullopt}})),
                    makeConstantExpr(
                        make_array_vector<StringView>({{StringView("6")}})),
                    makeConstantExpr(make_array_vector<Timestamp>(
                        {{Timestamp(123'456, 123'000)}})),
                    makeConstantExpr(make_array_vector<int32_t>({{8035}}, DATE())),
                    makeConstantExpr(make_array_vector<int64_t>(
                        {{54 * 1000}}, INTERVAL_DAY_TIME())),
                    makeConstantExpr(make_array_vector<int64_t>({{}})),
                    // Nested array: [[1, 2, 3], [4, 5]]
                    makeConstantExpr(make_array_vector(
                        {0}, make_array_vector<int64_t>({{1, 2, 3}, {4, 5}}))),
                };
                std::vector<std::string> names(expressions.size());
                for (auto i = 0; i < names.size(); ++i) {
                    names[i] = fmt::format("e{}", i);
                }
                return std::make_shared<core::ProjectNode>(
                    id, std::move(names), std::move(expressions), input);
            })
            .planNode();
    assertPlanConversion(
        plan,
        "SELECT array[true, null], array[0, null], array[1, null], "
        "array[2, null], array[3, null], array[4.4, null], array[5.5, null], "
        "array['6'],"
        "array['1970-01-02T10:17:36.000123000'::TIMESTAMP],"
        "array['1992-01-01'::DATE],"
        "array[INTERVAL 54 MILLISECONDS], "
        "array[], array[array[1,2,3], array[4,5]]");
}

TEST_F(PolluxSubstraitRoundTripTest, dateType) {
    auto a = make_flat_vector<int32_t>({0, 1});
    auto b = make_flat_vector<double_t>({0.3, 0.4});
    auto c = make_flat_vector<int32_t>({8036, 8035}, DATE());

    auto vectors = make_row_vector({"a", "b", "c"}, {a, b, c});
    createDuckDbTable({vectors});

    auto plan = PlanBuilder()
            .values({vectors})
            .filter({"c > DATE '1992-01-01'"})
            .planNode();
    assertPlanConversion(plan, "SELECT * FROM tmp WHERE c > DATE '1992-01-01'");
}

int main(int argc, char **argv) {
    testing::InitGoogleTest(&argc, argv);
    melon::Init init{&argc, &argv, false};
    return RUN_ALL_TESTS();
}
